# -*-coding:utf-8-*-
from mmdetzh.apis import init_detector, inference_detector, show_result_pyplot
import cv2
import os


class InferenceDetect:
    def __init__(self, img=None):
        self.config_file = '/home/ps/zihao/Detect_interf/resources/utils/' \
                           'Detection/source/cascade_rcnn_r50_fpn_1x_coco.py'
        self.checkpoint_file = '/home/ps/zihao/Detect_interf/resources/utils/Detection//source/epoch_9.pth'
        self.img = img
        self.model = init_detector(self.config_file, self.checkpoint_file, device='cuda:0')

    def image_detect(self, img=None):
        if img is None:
            if self.img is None:
                return
            img = self.img
        # print(img)
        result = inference_detector(self.model, img)
        detect_img = show_result_pyplot(self.model, img, result, score_thr=0.7)
        # print(type(detect_img))
        return detect_img, result

    @staticmethod
    def cv_show(name, img):
        cv2.imshow(f'{name}', img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    # def series_image_detect(self):
    #     img_list = [os.path.abspath(os.path.join(img_dir, img)) for img in os.listdir(img_dir)]
    #     result = inference_detector(model, img_list)
    #     all_bbox = []
    #     for one in result:
    #         bbox = np.concatenate((one[0], one[1]), axis=0).tolist()
    #         all_bbox.append(bbox)
    #     show_result_pyplot(model, img, result)


if __name__ == '__main__':
    img = '../../images/18.png'
    id = InferenceDetect(img)
    image = id.image_detect()
    # id.cv_show('haha', image)
